import json
import sys
import warnings
import numpy as np
import re
import os
import torch
import matplotlib.pyplot as plt

def DGI_reconstruction(y, patterns, batch_size, num_patterns, img_h, img_w, g_factors=0.5):
    patterns = torch.reshape(patterns, [num_patterns, img_h * img_w])
    y = torch.reshape(y, [batch_size, num_patterns, 1])

    compute1 = torch.transpose(
        (patterns - torch.matmul(torch.ones([num_patterns, 1]), torch.mean(patterns, dim=0).unsqueeze(0))), dim0=0,
        dim1=1)
    compute2 = torch.sum(patterns, dim=1)

    gammer = g_factors * torch.mean(y, dim=(1, 2)) / torch.mean(compute2)

    temp = (gammer.unsqueeze(-1) * compute2).unsqueeze(-1)

    # DGI process
    DGI = torch.matmul(compute1, y - temp).squeeze(-1)
    DGI = torch.reshape(DGI, [batch_size, 1, img_h, img_w])
    mean_DGI = torch.mean(DGI)
    std_DGI = torch.std(DGI)
    DGI = (DGI - mean_DGI) / std_DGI

    return DGI

# 获取输入参数
json_str = sys.argv[1]
try:
    file_list = json.loads(json_str)
    print(f"Received JSON string: {json_str}")
    print(f"Parsed file list: {file_list}")
except json.JSONDecodeError as e:
    print(f"Failed to parse JSON string: {e}")
    sys.exit(1)

pulse_count = int(sys.argv[2])
dir_name = sys.argv[3]
picture_name = sys.argv[4]
pt_path = sys.argv[5]

print(f"Pulse count: {pulse_count}")
print(f"Output directory: {dir_name}")
print(f"Picture name: {picture_name}")

y_np = []

# 循环计算不同数量的脉冲差值
for photo_num in [pulse_count]:
    accmulate = []  # 放置11个文本的脉冲差平均值
    
    for file in file_list:
        # 指定列索引
        column_index_1 = 1  # 第二列（标记）
        column_index_3 = 3  # 第四列（脉冲数）

        # 初始化数组
        mar_array = []
        pulse_array = []  # 存储第四列数据（脉冲数）

        # 读取文件
        with open(file, 'r', encoding='utf-8') as f:
            data_started = False
            for line in f:
                # 查找数据起始行
                if not data_started:
                    if line.startswith("record#"):
                        data_started = True
                    continue
                
                # 分割不同列的数据
                if data_started and line.strip():
                    row = re.split(r'\s+', line.strip())
                
                if len(row) > column_index_1:
                    mar_array.append(row[column_index_1])
                if len(row) > column_index_3:
                    pulse_array.append(row[column_index_3])

        # 将数据转换为NumPy数组
        MAR = np.array(mar_array)
        PULSES = np.array(pulse_array).astype(np.int64)  # 使用更大范围的整数类型
        
        # 清除MAR,PULSES中脉冲数为1时的数据
        pdel = []
        mdel = []
        for item in range(len(MAR)):
            if PULSES[item] == 1:
                pdel.append(item)
                mdel.append(item)
        PULSES = np.delete(PULSES, pdel)
        MAR = np.delete(MAR, mdel)
        
        # 获取MAR标记的位置
        indices_mar = np.where(MAR == 'MAR')[0]  # 注意：文件中的标记实际为'CHN'
        
        first_pulse_per_image = []
        # 建立循环，使用列表记录每次MAR后脉冲
        a = 0  # a作为索引定位到存储每次MAR后脉冲所对应的列表
        for index in indices_mar:
            # 将每个MAR后指定数量的脉冲差值放入一个列表
            first_pulse_per_image.append([])
            for i in range(1, photo_num + 1):
                real_index = index + 1
                if PULSES[real_index] - PULSES[index] == 0:  # 检测前一个脉冲与后一个脉冲重合的情况
                    index += 1
                    real_index = index + 1
                    first_pulse_per_image[a].append(PULSES[real_index] - PULSES[index])
                else:  # 正常情况
                    first_pulse_per_image[a].append(PULSES[real_index] - PULSES[index])
                index += 1
            a += 1
        
        list1 = []  # 记录每次脉冲差的平均值
        for line in first_pulse_per_image:
            mean = np.mean(line)
            list1.append(mean)
            mean_real = np.mean(list1)
        
        for item in list1:
            accmulate.append(1 / item)
    y_np = accmulate

# 转换为PyTorch张量
y_torch = torch.tensor(y_np, dtype=torch.float32)

# 重塑张量
y_torch = y_torch.reshape([1, 4096, 1, 1])

# 如果是打包后的可执行文件，使用 sys._MEIPASS 获取临时目录
if getattr(sys, 'frozen', False):
    current_dir = sys._MEIPASS
else:
    current_dir = os.path.dirname(os.path.abspath(__file__))

# 构建 patterns 文件的完整路径
patterns_path = os.path.join(current_dir, "patterns_hadamard_64.pt")
patterns_path = pt_path
# 加载 patterns
try:
    patterns = torch.load(patterns_path, map_location=torch.device('cpu'))
    print(f"Patterns loaded from: {patterns_path}")
except FileNotFoundError:
    print(f"Error: File not found at path: {patterns_path}")
    sys.exit(1)

DGI = DGI_reconstruction(y_torch, patterns, 1, patterns.shape[0], patterns.shape[2], patterns.shape[3], 0.5)

# 保存图像
output_file_name = os.path.join(dir_name, picture_name)
plt.figure(figsize=(15, 6))
plt.imshow(DGI.squeeze().numpy(), cmap="gray")  # 显示灰度图像
plt.title(f'Reconstructed Image using DGI with {photo_num} Photo')
plt.axis('off')
plt.savefig(output_file_name, bbox_inches='tight')
plt.close()

print(f"Reconstructed image saved as {output_file_name}")